import os
import cv2
import numpy as np
import pickle
import tensorflow as tf
from tensorflow.keras.applications.vgg16 import VGG16, preprocess_input
from tensorflow.keras.preprocessing import image
import faiss
import glob

# 调整图像大小为64x64以减少内存使用
def process_image(img_path, target_size=(64, 64)):
    try:
        img = cv2.imread(img_path)
        if img is None:
            return None
        img = cv2.resize(img, target_size)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = np.expand_dims(img, axis=0)
        img = preprocess_input(img)
        return img
    except Exception as e:
        print(f"处理图像 {img_path} 时出错: {str(e)}")
        return None

# 创建特征和标签数据
def createXY(train_folder=None, dest_folder='.', method='vgg16'):
    print(f"开始创建特征和标签...")
    
    # 定义可能的数据路径
    possible_paths = [
        train_folder,
        os.path.join('.', 'data', 'train'),
        os.path.join('.', 'cat_dog_data', 'cat_dog_data', 'data', 'train'),
        os.path.join('.', 'cat_dog_data', 'data', 'train'),
        os.path.join('.', 'train')
    ]
    
    # 寻找有效的数据路径
    train_path = None
    for path in possible_paths:
        if path and os.path.exists(path):
            # 检查目录结构
            cat_path = os.path.join(path, 'cat')
            dog_path = os.path.join(path, 'dog')
            if os.path.exists(cat_path) and os.path.exists(dog_path):
                train_path = path
                print(f"找到数据集路径: {train_path}")
                break
    
    if not train_path:
        raise ValueError(f"未找到有效的训练数据集路径，请检查数据是否正确放置。尝试过的路径: {possible_paths}")
    
    # 收集图像文件
    cat_files = glob.glob(os.path.join(train_path, 'cat', '*.jpg'))[:1000]  # 限制数量以减少内存使用
    dog_files = glob.glob(os.path.join(train_path, 'dog', '*.jpg'))[:1000]
    
    if not cat_files and not dog_files:
        # 尝试另一种可能的目录结构
        cat_files = glob.glob(os.path.join(train_path, 'cat*.jpg'))[:1000]
        dog_files = glob.glob(os.path.join(train_path, 'dog*.jpg'))[:1000]
    
    if not cat_files and not dog_files:
        raise ValueError(f"在 {train_path} 中未找到任何猫或狗的图像文件")
    
    print(f"找到 {len(cat_files)} 个猫图像和 {len(dog_files)} 个狗图像")
    
    # 采样以减少内存使用
    sample_ratio = 0.2  # 只使用20%的数据
    cat_sample_size = max(1, int(len(cat_files) * sample_ratio))
    dog_sample_size = max(1, int(len(dog_files) * sample_ratio))
    
    cat_files = np.random.choice(cat_files, cat_sample_size, replace=False)
    dog_files = np.random.choice(dog_files, dog_sample_size, replace=False)
    
    print(f"采样后: {len(cat_files)} 个猫图像和 {len(dog_files)} 个狗图像")
    
    # 加载模型
    if method == 'vgg16':
        model = VGG16(weights='imagenet', include_top=False)
    else:
        raise ValueError(f"不支持的特征提取方法: {method}")
    
    # 处理图像并提取特征
    X = []
    y = []
    
    for img_path in cat_files:
        img = process_image(img_path)
        if img is not None:
            features = model.predict(img)
            X.append(features)
            y.append(0)  # 猫为0
    
    for img_path in dog_files:
        img = process_image(img_path)
        if img is not None:
            features = model.predict(img)
            X.append(features)
            y.append(1)  # 狗为1
    
    if not X:
        raise ValueError("未能提取任何有效特征，请检查图像质量")
    
    # 转换为numpy数组
    X = np.array(X)
    y = np.array(y)
    
    # 创建目标文件夹（如果不存在）
    os.makedirs(dest_folder, exist_ok=True)
    
    # 保存数据
    with open(os.path.join(dest_folder, 'X.pkl'), 'wb') as f:
        pickle.dump(X, f)
    with open(os.path.join(dest_folder, 'y.pkl'), 'wb') as f:
        pickle.dump(y, f)
    
    print(f"特征和标签创建完成！形状: X={X.shape}, y={y.shape}")
    return X, y

# 安全的L2归一化函数
def safe_normalize_L2(X):
    try:
        # 检查X是否为空
        if X is None or len(X) == 0:
            print("警告: 输入数组为空，跳过归一化")
            return X
            
        # 确保X是numpy数组
        if not isinstance(X, np.ndarray):
            X = np.array(X)
            
        # 检查维度并转换为float32类型
        if len(X.shape) == 4:
            # 对于4D数组，将其展平为2D
            original_shape = X.shape
            X_flat = X.reshape(X.shape[0], -1).astype('float32')
            
            # 尝试使用faiss进行归一化
            try:
                faiss.normalize_L2(X_flat)
                # 恢复原始形状
                X_normalized = X_flat.reshape(original_shape)
                return X_normalized.astype('float16')  # 转换为float16以节省内存
            except Exception as e:
                print(f"faiss归一化失败: {str(e)}，使用numpy备选方法")
                # 使用numpy备选方法
                norms = np.linalg.norm(X_flat, axis=1, keepdims=True)
                norms[norms == 0] = 1  # 避免除零错误
                X_normalized = X_flat / norms
                return X_normalized.reshape(original_shape).astype('float16')
        elif len(X.shape) == 2:
            # 对于2D数组，直接使用faiss
            X_float32 = X.astype('float32')
            faiss.normalize_L2(X_float32)
            return X_float32.astype('float16')
        else:
            print(f"警告: 不支持的数组维度 {X.shape}，跳过归一化")
            return X
    except Exception as e:
        print(f"归一化过程中发生错误: {str(e)}")
        return X

# 展平特征函数
def flatten_features(X):
    if X is None or len(X) == 0:
        return X
    
    # 确保X是numpy数组
    if not isinstance(X, np.ndarray):
        X = np.array(X)
        
    # 检查维度
    if len(X.shape) == 4:
        # 对于4D数组 (样本数, 高度, 宽度, 通道)
        return X.reshape(X.shape[0], -1)
    elif len(X.shape) == 2:
        # 已经是2D数组，直接返回
        return X
    else:
        print(f"警告: 不支持的特征维度 {X.shape}")
        return X

# 准备特征函数
def prepare_features(X):
    # 展平特征
    X_flat = flatten_features(X)
    
    # 确保是float32类型
    if X_flat is not None and isinstance(X_flat, np.ndarray):
        X_flat = X_flat.astype('float32')
        
        # 执行L2归一化
        X_normalized = safe_normalize_L2(X_flat)
        
        # 转换为float16以节省内存
        if X_normalized is not None and isinstance(X_normalized, np.ndarray):
            X_normalized = X_normalized.astype('float16')
            
        return X_normalized
    
    return X_flat

# 加载数据函数
def load_data(file_path):
    try:
        with open(file_path, 'rb') as f:
            data = pickle.load(f)
        return data
    except Exception as e:
        print(f"加载数据文件 {file_path} 时出错: {str(e)}")
        return None

# 保存数据函数
def save_data(data, file_path):
    try:
        # 确保父目录存在
        os.makedirs(os.path.dirname(os.path.abspath(file_path)), exist_ok=True)
        with open(file_path, 'wb') as f:
            pickle.dump(data, f)
        print(f"数据已保存到 {file_path}")
    except Exception as e:
        print(f"保存数据文件 {file_path} 时出错: {str(e)}")